from flask import Blueprint, request, jsonify, current_app
import logging, os, json
from utils import common
from concurrent.futures import ThreadPoolExecutor
executor = ThreadPoolExecutor(10)

from ._impl import predict

sd_bp = Blueprint('sd', __name__, url_prefix='/sd')

@sd_bp.route('jobs', methods=['GET'])
def jobs():
    params = dict(request.args)
    result =  current_app.config['DB'].get_sd_jobs(params)

    ## 变成字典
    rdata = list()
    for item in result:
        rdata.append(
            {
                'id':item[0],
                'prompt': item[4],
                'img_url': item[5],
                'params': json.loads(item[6]),
            }
        )

    return jsonify({'code':0, 'data':rdata})


@sd_bp.route('models/list', methods=['GET'])
def models():
    params = dict(request.args)
    result =  current_app.config['DB'].get_sd_models(params)

    ## 变成字典
    rdata = list()
    for item in result:
        rdata.append(
            {
                'name': item[0]
            }
        )

    return jsonify({'code':0, 'data':rdata})

import random
import string
import datetime
 
def generate_random_code(length):
   characters = string.ascii_letters + string.digits
   random_code = ''.join(random.choice(characters) for _ in range(length))
   return random_code


@sd_bp.route('image', methods=['POST'])
def image():
    params = request.json
    
    ## 组装参数
    data = dict()
    data['task_id'] = generate_random_code(15)
    data['user_id'] = params.get('user_id')
    data['prompt'] = params.get('prompt')
    data['type'] = params.get('type')
    data['params'] = json.dumps(params, ensure_ascii=False)
    data['progress'] = 0
    data['created_at'] = datetime.datetime.now()
    data['img_url'] = ''
    data['publish'] = 0
    data['err_msg'] = ''
    data['progress_id'] = None
    

    result =  current_app.config['DB'].create_sd_job(data)
    ## 执行任务
    executor.submit(predict, data)

    return jsonify({'code':0})



@sd_bp.route('/remove', methods=['POST'])
def remove():
    params = request.json
    
    try:
        result =  current_app.config['DB'].remove_sd_job(params)
        return jsonify({'code':0})
    except Exception as e:
        return jsonify({'code': 500, 'messsage':str(e)})


@sd_bp.route('/publish', methods=['POST'])
def publish():
    params = request.json
    
    try:
        result =  current_app.config['DB'].publish_sd_job(params)
        return jsonify({'code':0})
    except Exception as e:
        return jsonify({'code': 500, 'messsage':str(e)})

@sd_bp.record
def log(app):
    logging.info('register sd blueprint')
